function [model infos] = trust_region_comparison(fun_set, fun_obj, data, model, params)
% Trust-region optimization on Riemannian manifolds
%
% This code has been adapted from the RTR - Riemannian Trust-Region toolbox
% (P.-A. Absil, C. G. Baker, K. A. Gallivan)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% set default parameters
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

t_begin = tic();
params = default_params(params);

user_tol = params.tol;
user_vtol = params.vtol;
user_grad_tol = params.grad_tol;
max_iter_tr = params.max_iter_tr;
verb = params.verb;
data_ts = params.data_ts; % Compute predictions at each iteration if asked


p = model.p;
Delta_bar = 10*p;   % Maximum trust-region radius
Delta0 = p;         % Initial trust-region radius
verbosity = 1;
min_inner = 0;
max_inner = intmax;
min_outer = 0; 3;  % Relax this for computing path regularization
max_outer = max_iter_tr;
epsilon = user_tol; % Outer Convergence tolerance (1e-6)
vepsilon = user_vtol;
kappa = 0.1; % Inner kappa convergence tolerance (0.1)
theta = 1;  % Inner theta convergence tolerance
rho_prime = 0.1;  %Accept/reject auxeter

% ***** Initializations *****

% initialize counters/sentinals
% allocate storage for dist, counters
k          = 0;  % counter for outer (TR) iteration.
stop_outer = 0;  % stopping criterion for TR.

% initialize solution and companion measures:
cache = struct;
[fx, cache] = feval(fun_set, 'f', fun_obj, model, [], [], data, params, cache); 
[fgradx, cache] = feval(fun_set, 'grad_f', fun_obj, model, [], [], data, params, cache);

% initialize trust-region radius
Delta = Delta0;

% **********************
% ** Start of TR loop **
% **********************

costs = zeros(max_outer+1,1);
costs(1) = fx;

iter_time = zeros(max_outer+1,1);
iter_time(1) = 0;

% If interested in computing recovery
if params.compute_predictions,
    preds_test = partXY((model.U*model.B)', model.V', data_ts.rows, data_ts.cols, data_ts.nentries)';
    errors_test = preds_test - data_ts.entries;
    cost_test = (errors_test'*errors_test);
    test_costs(1) = cost_test;
end


if verb
  fprintf('[%0.4d] cost = %.5e\n',0,costs(1));
end

while stop_outer==0,

    % update counter
    k = k+1;
    params.iter = k;

    % *************************
    % ** Begin TR Subproblem **
    % *************************
    fx_old=fx;
    
    % solve TR subproblem
    [eta,inner_it,stop_inner,cache] = tCG(fun_set, model, fgradx, Delta, theta, kappa, min_inner, max_inner, fun_obj, data, params, cache);
    if stop_inner == 1,
        rateConvergence = 'Negative curvature';
    elseif stop_inner == 2,
        rateConvergence = 'Exceeded trust region';
    elseif stop_inner == 3,
        rateConvergence = 'Linear';
    elseif stop_inner == 4,
        rateConvergence = 'Super-linear';
    end
    
    
    % do we accept the proposed solution or not?
    % compute the Hessian at the proposal
    [Heta, cache] = feval(fun_set,'hessian', fun_obj, model, eta, [], data, params, cache);
    
    % compute the retraction of the proposal (eta)
    [model_prop, cache]  = feval(fun_set,'retraction', fun_obj, model, eta, 1, data, params, cache);

    % compute function value of the proposal
    [fx_prop, cache] = feval(fun_set, 'f', fun_obj, model_prop, [], [], data, params, cache);

    % check the performance of the quadratic model
    rhonum = fx-fx_prop;
    rhoden = -feval(fun_set, 'metric', fun_obj, model, fgradx, eta, data, params, cache) - 0.5*feval(fun_set, 'metric', fun_obj, model, Heta, eta, data, params, cache); 
    rho =   rhonum  / rhoden;

    % choose new TR radius based on performance
    if rho < 1/4
        Delta = 1/4*Delta;
    elseif rho > 3/4 && (stop_inner == 2 || stop_inner == 1),
        Delta = min(2*Delta, Delta_bar);
    end

    % choose new iterate based on performance
    if rho > rho_prime,
        accept = true;
        model = model_prop;
        fx = fx_prop;
        [fgradx,cache] = feval(fun_set, 'grad_f', fun_obj, model, [], [], data, params, cache); %fns.fgrad(x);
    else
        accept = false;
    end
    
    norm_grad = sqrt(feval(fun_set, 'metric', fun_obj, model, fgradx, fgradx, data, params, cache));

    % If interested in computing recovery
    if params.compute_predictions,
        preds_test = partXY((model.U*model.B)', model.V', data_ts.rows, data_ts.cols, data_ts.nentries)';
        errors_test = preds_test - data_ts.entries;
        cost_test = (errors_test'*errors_test);
        test_costs(k+1) = cost_test;
    end

    
    % ** Testing for Stop Criteria
    % min_outer is the minimum number of inner iterations
    % before we can exit. this gives randomization a chance to
    % escape a saddle point.

%     if accept && k > min_outer && (abs(rhonum) < epsilon || abs(rhonum)/fx_old < vepsilon),
%     if k > min_outer && (abs(rhonum) < epsilon || abs(rhonum)/fx_old < vepsilon),
%     if k > min_outer && norm_grad < user_grad_tol,
%     if k > min_outer && (norm_grad < user_grad_tol || abs(rhonum) < epsilon || abs(rhonum)/fx_old < vepsilon),
    if k > min_outer && fx < user_tol,
   
        stop_outer = 1;
    end

    if k >= max_outer,
        if (verbosity > 0),
            fprintf('Maximum number of iterations (%d) is reached \n',k);
        end
        stop_outer = 1;
    end

    iter_time(k+1) = toc(t_begin);
    costs(k+1) = fx;
    
    if verb
        %%% Bark only when asked
        fprintf('[%0.4d] cost = %.5e  gradnorm = %.5e  inner iter. = %0.3i  radius = %.3e  %s \n',k,costs(k+1), norm_grad, inner_it, Delta, rateConvergence);
    end
end  % of TR loop (counter: k)

infos.cost = fx;
infos.costs = costs(1:k+1);
infos.iter_time = iter_time(1:k+1);
infos.iter = k+1;
if params.compute_predictions,
    infos.test_error = test_costs(1:k+1);
    infos.test_error = infos.test_error';
end

return;



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%% truncated CG %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [eta, inner_it, stop_tCG, cache] = tCG(fun_set, x, grad, Delta, theta, kappa, min_inner, max_inner, fun_obj, data, params, cache)
% tCG - Truncated (Steihaug-Toint) Conjugate-Gradient
% minimize <eta,grad> + .5*<eta,Hess(eta)>
% subject to <eta,eta> <= Delta^2

% all terms involving the trust-region radius will utilize an inner product
% w.r.t. the preconditioner; this is because the iterates grow in
% length w.r.t. the preconditioner, guaranteeing that we will not
% re-enter the trust-region
%
% the following recurrences for Prec-based norms and inner
% products come from CGT2000, pg. 205, first edition
% below, P is the preconditioner
%
% <eta_k,P*delta_k> = beta_k-1 * ( <eta_k-1,P*delta_k-1> + alpha_k-1 |delta_k-1|^2_P )
% |delta_k|^2_P = <r_k,z_k> + beta_k-1^2 |delta_k-1|^2_P
%
% therefore, we need to keep track of
% 1)   |delta_k|^2_P
% 2)   <eta_k,P*delta_k> = <eta_k,delta_k>_P
% 3)   |eta_k  |^2_P
%
% initial values are given by:
%    |delta_0|_P = <r,z>
%    |eta_0|_P   = 0
%    <eta_0,delta_0>_P = 0
% because we take eta_0 = 0


[eta, cache] = feval(fun_set, 'zero_dir', fun_obj, x, [], [], data, params, cache);
% eta = 0*grad;

r = grad;
e_Pe = 0;

r_r = feval(fun_set, 'metric', fun_obj, x, r, r, data, params, cache); %fns.g(x,r,r);
norm_r = sqrt(r_r);
norm_r0 = norm_r;
z=r;
z_r = r_r;
d_Pd = z_r;

% initial search direction
[delta, cache] = feval(fun_set, 'scale_dir', fun_obj, [], z, -1, data, params, cache);
% delta = -z;    

e_Pd = 0;

% pre-assume termination b/c j == end
stop_tCG = 5;

% begin inner/tCG loop
j = 0;
for j = 1 : max_inner,
    [Hxd, cache] = feval(fun_set, 'hessian', fun_obj, x, delta, [], data, params, cache); 
    
    % compute curvature
    [d_Hd, cache] = feval(fun_set, 'metric', fun_obj, x, delta, Hxd, data, params, cache); 

    alpha = z_r/d_Hd;
    % <neweta,neweta>_P = <eta,eta>_P + 2*alpha*<eta,delta>_P + alpha*alpha*<delta,delta>_P
    e_Pe_new = e_Pe + 2.0*alpha*e_Pd + alpha*alpha*d_Pd;

    if d_Hd <= 0 || e_Pe_new >= Delta^2,
        % want
        %  ee = <eta,eta>_prec,x
        %  ed = <eta,delta>_prec,x
        %  dd = <delta,delta>_prec,x
        tau = (-e_Pd + sqrt(e_Pd*e_Pd + d_Pd*(Delta^2-e_Pe))) / d_Pd;
        
        [eta, cache] = feval(fun_set, 'add_dir', fun_obj, eta, delta, tau, data, params, cache); 
        % eta = eta + tau*delta; CHECK
        
        if d_Hd <= 0,
            stop_tCG = 1;     % negative curvature
        else
            stop_tCG = 2;     % exceeded trust region
        end
        break;
    end

    % no negative curvature and eta_prop inside TR: accept it
    e_Pe = e_Pe_new;
    
    [eta, cache] = feval(fun_set, 'add_dir', fun_obj, eta, delta, alpha, data, params, cache); 
    % eta = eta + alpha*delta;

    % update the residual
    [r, cache] = feval(fun_set, 'add_dir', fun_obj, r, Hxd, alpha, data, params, cache); 
    % r = r + alpha*Hxd;
   
    % re-tangentalize r
    [r, cache] = feval(fun_set, 'project', fun_obj, x, r, [], data, params, cache); 

    % compute new norm of r
    r_r = feval(fun_set, 'metric', fun_obj, x, r, r, data, params, cache); 
    norm_r = sqrt(r_r);

    % check kappa/theta stopping criterion
   if j >= min_inner && norm_r <= norm_r0*min(norm_r0^theta,kappa)       
        % residual is small enough to quit
        if kappa < norm_r0^theta,
            stop_tCG = 3;  % linear convergence
        else
            stop_tCG = 4;  % superlinear convergence
        end
        break;
    end

    % precondition the residual
    z=r;
    
    % save the old z'*r
    zold_rold = z_r;
    % compute new z'*r
    z_r = r_r;

    % compute new search direction
    beta = z_r/zold_rold;
    
    [delta, cache] = feval(fun_set,'scale_dir', fun_obj, [], delta, beta, data, params, cache);
    [delta, cache] = feval(fun_set, 'add_dir', fun_obj, delta, z, -1, data, params, cache);
    % delta = -z + beta*delta; CHECK
    
    % update new P-norms and P-dots
    e_Pd = beta*(e_Pd + alpha*d_Pd);
    d_Pd = z_r + beta*beta*d_Pd;

end  % of tCG loop
inner_it = j;
return;

